# DecisionTree Classifier
from src.models.supmodel import SupModel
from sklearn.tree import DecisionTreeClassifier

class DecisionTree(SupModel):
    """
    @brief 决策树分类器（支持决策树桩）
    @details 基于sklearn的DecisionTreeClassifier实现，支持max_depth参数和sample_weight。
    """
    def __init__(self, max_depth=1):
        """
        @brief 初始化决策树分类器
        @param max_depth 决策树最大深度，默认为1（决策树桩）
        """
        self.max_depth = max_depth
        self.model = DecisionTreeClassifier(max_depth=self.max_depth)
        self.name = f"DT_{self.max_depth}"

    def fit(self, X_train, y_train, sample_weight=None):
        """
        @brief 训练决策树模型
        @param X_train 训练特征
        @param y_train 训练标签
        @param sample_weight 样本权重（可选）
        """
        if sample_weight is not None:
            self.model.fit(X_train, y_train, sample_weight=sample_weight)
        else:
            self.model.fit(X_train, y_train)

    def predict(self, X):
        """
        @brief 对输入数据进行预测
        @param X 测试特征
        @return 预测标签
        """
        return self.model.predict(X)
